from kubernetes import client, config

def get_pods(namespace, label_selector, node_name):
    """获取符合条件的Pods"""
    v1 = client.CoreV1Api()
    pods = v1.list_namespaced_pod(namespace, label_selector=label_selector, field_selector=f'spec.nodeName={node_name}')
    return pods.items

def patch_pod_labels(pod, new_labels):
    """更新Pod的标签，并返回旧标签以便回退"""
    v1 = client.CoreV1Api()
    old_labels = pod.metadata.labels  # 保存旧标签

    patch = {"metadata": {"labels": new_labels}}
    
    try:
        v1.patch_namespaced_pod(name=pod.metadata.name, namespace=pod.metadata.namespace, body=patch)
        print(f"Successfully patched pod {pod.metadata.name}.")
        return old_labels  # 返回旧标签
    except Exception as e:
        print(f"Error patching pod {pod.metadata.name}: {e}")
        return None  # 返回None表示失败

def rollback_pod_labels(pod, old_labels):
    """回退到旧标签"""
    if old_labels is None:
        print(f"No old labels to roll back for pod {pod.metadata.name}.")
        return

    v1 = client.CoreV1Api()
    patch = {"metadata": {"labels": old_labels}}

    try:
        v1.patch_namespaced_pod(name=pod.metadata.name, namespace=pod.metadata.namespace, body=patch)
        print(f"Rolled back pod {pod.metadata.name} to old labels.")
    except Exception as e:
        print(f"Error rolling back pod {pod.metadata.name}: {e}")

def main():
    # 加载 Kubernetes 配置
    config.load_kube_config()  # 如果在集群外部运行

    namespace = "default"
    label_selector = "app=inference"
    node_name = "cn-shenzhen.192.168.1.101"
    new_label = {"app": "inference-bak"}

    # 获取指定节点上符合条件的 Pods
    pods = get_pods(namespace, label_selector, node_name)

    # 用于保存每个 pod 的旧标签
    old_labels_mapping = {}

    # 更新每个 pod 的标签并保存旧标签
    for pod in pods:
        old_labels = patch_pod_labels(pod, new_label)
        old_labels_mapping[pod.metadata.name] = old_labels

    # 假设某种条件触发了回退，实际使用中你可以根据需求决定何时调用回退
    # 例如: 回退所有的 Pods 标签
    for pod in pods:
        rollback_pod_labels(pod, old_labels_mapping.get(pod.metadata.name))

if __name__ == "__main__":
    main()
